#include "kernel_operator.h"

constexpr uint32_t OUTPUT_DATA_TYPE_SIZE = 4;                           // 输出float的字节数
constexpr uint32_t BLOCK_SIZE = 32;                                     // 32B内存对齐基准
constexpr uint32_t TILE_ELEM_NUM = BLOCK_SIZE / OUTPUT_DATA_TYPE_SIZE;  // float的32B对齐元素数
constexpr uint32_t BUFFER_NUM = 2;                                      // 双缓冲

class KernelRange {
public:
    __aicore__ inline KernelRange() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR s, GM_ADDR z, uint32_t _totalLength,
                                uint32_t _blockLength, uint32_t _tileNum,
                                uint32_t _lastTileLength, uint32_t _formerLength,
                                uint32_t _formerTileNum, uint32_t _formerLastTileLength,
                                uint32_t _formerNum, uint32_t _tailLength,
                                uint32_t _tailTileNum, uint32_t _tailLastTileLength,
                                uint32_t _isEvenCore)
    {
        // 不论输入是什么类型，kernel侧始终使用float进行计算
        this->start = *reinterpret_cast<const __gm__ float*>(x);
        this->end = *reinterpret_cast<const __gm__ float*>(y);
        this->step = *reinterpret_cast<const __gm__ float*>(s);
        this->origSize = _totalLength;

        uint32_t coreId = AscendC::GetBlockIdx();
        uint32_t currentBlockLength = 0;

        if (_isEvenCore)
        {
            this->blockLength = _blockLength;
            this->tileNum = _tileNum;
            this->lastTileLength = _lastTileLength;
            this->blockOffset = blockLength * coreId;
            currentBlockLength = blockLength;
        }
        else
        {
            if (coreId < _formerNum)
            {
                this->blockLength = _formerLength;
                this->tileNum = _formerTileNum;
                this->lastTileLength = _formerLastTileLength;
                this->blockOffset = _formerLength * coreId;
                currentBlockLength = _formerLength;
            }
            else
            {
                this->blockLength = _tailLength;
                this->tileNum = _tailTileNum;
                this->lastTileLength = _tailLastTileLength;
                this->blockOffset = _formerLength * _formerNum + _tailLength * (coreId - _formerNum);
                currentBlockLength = _tailLength;
            }
        }
        zGm.SetGlobalBuffer(reinterpret_cast<__gm__ float*>(z) + this->blockOffset, currentBlockLength);
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, TILE_ELEM_NUM * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        if (origSize == 1)
        {
            AscendC::LocalTensor<float> tileOutput = outQueueZ.AllocTensor<float>();
            for (int32_t i = 0; i < TILE_ELEM_NUM; i++) tileOutput.SetValue(i, (i == 0) ? start : 0.0f);
            outQueueZ.EnQue(tileOutput);
            CopyOut(0);
            return;
        }
        for (int32_t tileIdx = 0; tileIdx < tileNum; tileIdx++)
        {
            Compute(tileIdx);
            CopyOut(tileIdx);
        }
    }

private:
    __aicore__ inline void CopyIn() {}
    
    __aicore__ inline void Compute(int32_t tileIdx)
    {
        AscendC::LocalTensor<float> tileOutput = outQueueZ.AllocTensor<float>();
        int32_t tileElemNum = TILE_ELEM_NUM;  // 8
        int32_t tileGlobalStart = blockOffset + tileIdx * tileElemNum;

        for (int32_t i = 0; i < tileElemNum; i++)
        {
            int32_t globalIdx = tileGlobalStart + i;
            if (globalIdx < origSize)
                tileOutput.SetValue(i, start + step * static_cast<float>(globalIdx));
            else
                tileOutput.SetValue(i, 0.0f);
        }
        outQueueZ.EnQue(tileOutput);
       
    }

    __aicore__ inline void CopyOut(int32_t tileIdx)
    {
        AscendC::LocalTensor<float> zLocal = outQueueZ.DeQue<float>();
        int32_t copyElemNum;
        if (tileIdx == tileNum - 1) copyElemNum = lastTileLength;
        else copyElemNum = TILE_ELEM_NUM;
        
        int32_t copyOffset = tileIdx * TILE_ELEM_NUM;
        if (copyOffset + copyElemNum > blockLength) copyElemNum = blockLength - copyOffset;
        
        AscendC::DataCopy(zGm[copyOffset], zLocal, copyElemNum);
        outQueueZ.FreeTensor(zLocal);

    }

private:
    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::GlobalTensor<float> zGm;

    uint32_t blockLength;    // 当前核心处理的float总元素数
    uint32_t tileNum;        // 当前核心的Tile总数（每个Tile 8个float元素）
    uint32_t lastTileLength; // 最后一个Tile的float元素数
    uint32_t blockOffset;    // 当前核心的全局float元素偏移
    uint32_t origSize;       // 原始有效元素数
    float start;
    float end;
    float step;
};

extern "C" __global__ __aicore__ void range_custom(GM_ADDR start, GM_ADDR end, GM_ADDR step, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    
    // 直接使用KernelRange类，不论输入类型，kernel侧始终使用float进行计算
    KernelRange op;
    op.Init(start, end, step, z, tiling_data.totalLength,
            tiling_data.blockLength, tiling_data.tileNum,
            tiling_data.lastTileLength, tiling_data.formerLength,
            tiling_data.formerTileNum, tiling_data.formerLastTileLength,
            tiling_data.formerNum, tiling_data.tailLength,
            tiling_data.tailTileNum, tiling_data.tailLastTileLength,
            tiling_data.isEvenCore);
    op.Process();
}
